import json  # 导入 JSON 模块，用于处理 JSON 数据
from collections import defaultdict


# ------------------------------
# Sink算子
# ------------------------------
class FileSink:
    def __init__(self, file_path):
        self.file_path = file_path

    def __call__(self, data):
        """将窗口统计结果写入文件"""
        # with self.lock:  # 使用锁确保线程安全
        with open(self.file_path, 'a') as f:  # 以追加模式打开文件
            window_time = data['window_end']  # 获取窗口结束时间
            results = json.dumps((data['key'], data['count']), ensure_ascii=False)  # 将结果转换为 JSON 字符串，确保支持非 ASCII 字符
            if data.get('is_retract', False):
                f.write(f"Retract: [{window_time}] {results}\n")
            else:
                f.write(f"Output: [{window_time}] {results}\n")

            # # results = json.dumps((data['key'],data['count']), ensure_ascii=False)  # 将结果转换为 JSON 字符串，确保支持非 ASCII 字符
            # f.write(f"[{window_time}] {results}\n")  # 将格式化后的结果写入文件


# ------------------------------
# Map 算子
# ------------------------------

class MapOperator:
    def __init__(self, func=lambda x: {'word': x['word'], 'count': x['count'], 'timestamp': x['timestamp'],'_event_id': x['_event_id'],'_watermark':x['_watermark']}):
        """初始化 Map 算子"""
        self.func = func  # 设置映射函数，默认将记录中的 `word`、`count` 和 `timestamp` 提取出来

    def __call__(self, record):
        """对记录执行映射操作"""
        return self.func(record)  # 调用映射函数处理记录并返回结果

# ------------------------------
# Reduce 算子
# ------------------------------
class ReduceOperator:
    """初始化 Reduce 算子"""
    def __init__(self, func):
        self.func = func

    """对记录执行聚合操作"""
    def __call__(self, records):
        return self.func(records)  # 调用分组函数处理记录并返回分组键
# ------------------------------
# KeyBy 算子
# ------------------------------
class KeyByOperator:
    def __init__(self, func=lambda x: x['word']):
        """初始化 KeyBy 算子"""
        self.func = func  # 设置分组函数，默认以 `word` 字段作为分组键

    def __call__(self, records):
        """对记录执行分组操作"""
        return self.func(records)  # 调用分组函数处理记录并返回分组键
# ------------------------------
# Retract 算子
# ------------------------------
class RetractOperator:
    def __init__(self, func, checkpoint_manager, node_name):
        self.func = func
        self.checkpoint_manager = checkpoint_manager
        self.node_name = node_name
        self.state = self.checkpoint_manager.restore_checkpoint(node_name) or defaultdict(int)

    def __call__(self, record):
        if record.get('is_retract', False):
            self.state[record['word']] -= record['count']
        else:
            self.state[record['word']] += record['count']
        # 保存状态到检查点
        self.checkpoint_manager.save_checkpoint(self.node_name, self.state)
        return {'word': record['word'], 'count': self.state[record['word']]}
# ------------------------------
# ExactlyOnce算子
# ------------------------------
class ExactlyOnceSink:
    def __init__(self, file_path='exactonle.txt'):
        self.file_path = file_path
        self.processed_ids = set()

    def __call__(self, data):
        if data['_event_id'] not in self.processed_ids:
            with open(self.file_path, 'a') as f:
                window_time = data['window_end']  # 获取窗口结束时间
                results = json.dumps((data['key'], data['count']), ensure_ascii=False)  # 将结果转换为 JSON 字符串，确保支持非 ASCII 字符
                f.write(f"[{window_time}] {results}\n")
            self.processed_ids.add(data['_event_id'])